# -*- coding: utf-8 -*-
"""
y ： 起报时间
x : 预报 lead time
对不同集合成员id求RMSE
"""

import pandas as pd
import os,glob,math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
from tqdm import tqdm
from global_land_mask import globe
from typlot.scripts.gsj_typhoon import tydat,see,count_rapidgrow,tydat_CMA,average_datetime,split_str_id,load_land_polygons,detect_landfall
from geopy.distance import geodesic
import matplotlib.ticker as ticker
import seaborn as sns
import xarray as xr
import warnings
warnings.filterwarnings("ignore")
from typlot.config.global_config import *
from geopy.distance import geodesic

ini_time_mode = ['00','12']
names = ['mojie_28','dusurui_16','gaemi_09','haikui_38','kangni_54','shantuo_44','saola_25','koinu_49']
# names = ['dusurui_16']
tynames,tyids = split_str_id(names)
draw_obs_opt = True
obs_baseline='land'  # ‘land’ 'RI'
RIstd = 7
tyrmse = {}
show = False
track_id = np.arange(1,52)

#  制作数据集
for ty,tyid in zip(tynames,tyids):  
    # 初始化
    directory  = os.path.join(global_ensdir,f'{ty}_{tyid}')
    dates_name = sorted(os.listdir(directory))
    dates_name = [i for i in dates_name if i[-2:] in ini_time_mode ]
    obs_path   = os.path.join( global_obsdir, f'{ty}_CMAobs.txt')
    tyobs = tydat_CMA(obs_path)
    variable = ['umax','pmin','dist']

    if not os.path.exists(os.path.join(basedir,f'{ty}_pmin_umax_dist.nc')):
        #### 计算最长的leadtime，确定后用来制作DataArray
        leadtime_list=[]
        # 遍历起报时间
        for date_name in tqdm(dates_name,total=len(dates_name),desc='Filtering max leadtime'):
            dir_date = os.path.join(directory,date_name)
            name_ensembles = os.listdir(dir_date)
            path_ensembles = [os.path.join(dir_date,i) for i in name_ensembles if i.startswith('TRACK')]
            # 遍历集合
            for path in path_ensembles:
                t_id = int(path.split('_')[-1])
                if t_id==0:
                    continue # ERA5 情况
                else:
                    tyens = tydat(path,RIstd)
                    ini_time = datetime.strptime(date_name, '%Y%m%d%H')
                    end_time = tyens.time[-1]
                    leadtime_list.append(end_time-ini_time)
        
        max_leadtime = max(leadtime_list)
        n_leadtime = np.ceil(max_leadtime / timedelta(hours=6)).astype(int)+1
        
        #### 创建DataArray  dims=['variable','start_time','lead_time']
        lead_time = 6*np.arange(n_leadtime)  # 6 for 6h/index.
        da = xr.DataArray(np.nan,dims=['variable','start_time','lead_time'],
                    coords={'variable':variable ,'start_time':dates_name ,'lead_time':lead_time})

        #### 开始填充da,逐个元素计算
        # 遍历start_time
        for date_name in tqdm(dates_name,total=len(dates_name),desc='Calcing RMSE on members'):
            dir_date = os.path.join(directory,date_name)
            name_ensembles = os.listdir(dir_date)
            path_ensembles = [os.path.join(dir_date,i) for i in name_ensembles if i.startswith('TRACK')]
            # 遍历 lead_time
            v_ens_list,p_ens_list = [],[]
            v_obs_list,p_obs_list = [],[]
            lat_ens_list,lat_obs_list = [],[]
            lon_ens_list,lon_obs_list = [],[]
            for i in lead_time:
                leadtime = datetime.strptime(date_name,'%Y%m%d%H') + timedelta(hours=int(i))
                # 遍历所有的成员，计算
                for path in path_ensembles:
                    t_id = int(path.split('_')[-1])
                    if t_id==0:
                        continue # 再分析 情况
                    else:
                        tyens = tydat(path,RIstd)
                        # 只留下leadtime同时存在于ens和obs的情况
                        if (leadtime in tyobs.time) and (leadtime in tyens.time):
                            v_ens_list.append(tyens.umax[tyens.time==leadtime][0])
                            v_obs_list.append(tyobs.umax[tyobs.time==leadtime][0])
                            p_ens_list.append(tyens.pmin[tyens.time==leadtime][0])
                            p_obs_list.append(tyobs.pmin[tyobs.time==leadtime][0])
                            lat_ens_list.append(tyens.lat[tyens.time==leadtime][0])
                            lat_obs_list.append(tyobs.lat[tyobs.time==leadtime][0])
                            lon_ens_list.append(tyens.lon[tyens.time==leadtime][0])
                            lon_obs_list.append(tyobs.lon[tyobs.time==leadtime][0])
                        else:
                            v_ens_list.append(np.nan)
                            v_obs_list.append(np.nan)
                            p_ens_list.append(np.nan)
                            p_obs_list.append(np.nan)
                            lat_ens_list.append(np.nan)
                            lat_obs_list.append(np.nan)
                            lon_ens_list.append(np.nan)
                            lon_obs_list.append(np.nan)
                v_ens,p_ens,v_obs,p_obs = np.array(v_ens_list),np.array(p_ens_list),np.array(v_obs_list),np.array(p_obs_list)
                lat_ens,lat_obs,lon_ens,lon_obs = np.array(lat_ens_list),np.array(lat_obs_list),np.array(lon_ens_list),np.array(lon_obs_list)
                dist_rmse_list = [geodesic((i,j),(k,l)).kilometers if not math.isnan(i*j*k*l) else np.nan for i,j,k,l in zip(lat_ens,lon_ens,lat_obs,lon_obs) ]
                
                #### 计算该leadtime的RMSE
                umax_rmse = np.sqrt(np.nanmean((v_ens-v_obs)**2))
                pmin_rmse = np.sqrt(np.nanmean((p_ens-p_obs)**2))
                dist_rmse = np.sqrt(np.nanmean((np.array(dist_rmse_list))**2))

                #### 补充DataArray
                da.loc['umax',date_name,i] = umax_rmse
                da.loc['pmin',date_name,i] = pmin_rmse
                da.loc['dist',date_name,i] = dist_rmse

        da.to_netcdf(os.path.join(basedir,f'{ty}_pmin_umax_dist.nc'))
    else :
        print('load existed dataarray')
        da = xr.open_dataarray(os.path.join(basedir,f'{ty}_pmin_umax_dist.nc'))



    def get_RI_init_time(tyobs, da):
        """
        根据tyobs的信息，找到观测中的RI时刻。
        但是观测RI往往不在严格的00，12h这种起报时间上。
        因此要把观测RI转换到最近的RI时刻上去。
        """
        from datetime import datetime
        index = count_rapidgrow(7, tyobs.umax, tyobs.time).astype(bool)
        obs_RI_moments = tyobs.time[index]
        # 将 da.start_time 从字符串转换为 datetime 对象
        da_start_times = [datetime.strptime(str(t), "%Y%m%d%H") for t in da.start_time.values]
        # 找到da中最近的RI起报时刻
        da_RI_moments = []
        for obs_time in obs_RI_moments:
            # 计算时间差
            time_diffs = np.array([(dt - obs_time).total_seconds() for dt in da_start_times])
            time_diffs[time_diffs < 0] = np.inf
            # 找到最小正数时间差的索引
            min_idx = np.argmin(time_diffs)
            # 添加对应的 start_time 字符串
            da_RI_moments.append(da.start_time.values[min_idx])
        return da_RI_moments

    def ceil_timedelta(td):
        """将 Timedelta 向上取整到天"""
        return pd.Timedelta(days=int(np.ceil(td.total_seconds() / 86400)))


    def gen_RI_scatter(tyobs,da):
        # 根据da来得到y和x轴的时间维度信息
        init_times = [pd.Timestamp(i+'00') for i in da.start_time.values]
        dy = init_times[1]-init_times[0]
        dx = pd.Timedelta(da.lead_time.values[1]-da.lead_time.values[0],'h')

        RI_ini_times = [pd.Timestamp(i+'00') for i in get_RI_init_time(tyobs, da)]
        x = []
        y = []
        for RI_ini_time in RI_ini_times:
            ini_time = RI_ini_time
            while ini_time>=init_times[0] and RI_ini_time-ini_time <= pd.Timedelta(da.lead_time.values[-1],'h'):
                x.append((RI_ini_time-ini_time))
                y.append(ini_time)
                ini_time -= dy
        # 对x进行处理
        x = [ceil_timedelta(xx) for xx in x]

        # 转换为index
        yindex = np.searchsorted(init_times,np.array(y))
        lead_time_timedelta = [pd.Timedelta(i,'h') for i in da.lead_time.values]
        xindex = np.searchsorted(lead_time_timedelta,np.array(x))
        return yindex,xindex


    for var in variable:
        # 1) 读取数据               
        data = da.loc[var].values.astype(float)               # 转 float，避免 int 溢出
        data = np.where(data > 10000, np.nan, data)  # 屏蔽异常值
        data = np.ma.masked_invalid(data)            # NaN → masked

        # 2) 坐标 
        members = da.lead_time.values  
        times   = da.start_time.values 

        # 🔧 将预报时效转换为天数，并筛选每天一次
        lead_days = members / 24  # 小时转换为天数
        day_indices = []
        day_labels = []

        for i, day in enumerate(lead_days):
            if day.is_integer():  # 只保留整天数
                day_indices.append(i)
                day_labels.append(f'{int(day)}d')

        # 如果没有整天数，则按24小时间隔筛选. 但是预报总是有整天数。
        if not day_indices:
            step = 24  # 每24小时取一个点
            day_indices = list(range(0, len(members), step))
            day_labels = [f'{int(members[i]/24)}d' for i in day_indices]
            
        # 3) 画图
        fig, ax = plt.subplots(figsize=(12, 8))

        # ⭐ 根据变量类型设置颜色范围
        if var == 'pmin':
            vmin, vmax = 0, 50
            cbar_label = 'Pmin RMSE (hPa)'
        elif var == 'umax':
            vmin, vmax = 0, 30
            cbar_label = 'Umax RMSE (m/s)'
        else:
            vmin, vmax = 0, 2500
            cbar_label = 'Geodesic Distance'

        # ✅ 在 imshow 中设置 vmin/vmax
        im = ax.imshow(
            data,                    
            aspect='auto',
            interpolation='none',
            cmap='viridis',
            origin='lower',
            vmin=vmin,  # ✅ 正确位置
            vmax=vmax   # ✅ 正确位置
        )

        # 🔧 交换 X、Y 轴的刻度和标签设置
        ax.set_yticks(np.arange(len(times)))
        ax.set_yticklabels(times, rotation=0, fontsize=8)

        # 🔧 Y轴只显示整天数的刻度
        ax.set_xticks(day_indices)
        ax.set_xticklabels(day_labels, fontsize=8)

        # 🔧 交换轴标签
        ax.set_title(f'{ty} — {var} RMSE Heatmap', fontsize=16, fontweight='bold')
        ax.set_ylabel('Forecast Initialization Time', fontsize=14)
        ax.set_xlabel('Forecast Lead Time', fontsize=14)

        # ⭐ 色标（不需要 vmin/vmax）
        cbar = fig.colorbar(im, ax=ax)
        cbar.set_label(cbar_label, fontsize=12, fontweight='bold')

        # scatter plot
        y,x = gen_RI_scatter(tyobs,da)
        ax.scatter(x,y,color='red')

        plt.tight_layout()

        if show:
            plt.show()
        else:
            plt.savefig(
                os.path.join(global_picdir,'RMSE' ,f'heatmap_{ty}_{var}_lead_start_days.png'),
                dpi=600
            )
            plt.close()